package com.filetransfer.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.filetransfer.model.dto.MessageDTO;
import com.filetransfer.service.DeviceService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket服务器
 */
@Slf4j
@Component
@ServerEndpoint("/ws/{deviceId}")
@RequiredArgsConstructor
public class WebSocketServer {

    /**
     * 已连接的会话
     */
    private static final Map<String, Session> SESSIONS = new ConcurrentHashMap<>();
    
    /**
     * JSON序列化工具
     */
    private static ObjectMapper objectMapper;
    
    /**
     * 设备服务
     */
    private static DeviceService deviceService;
    
    @PostConstruct
    public void init() {
        WebSocketServer.objectMapper = new ObjectMapper();
        log.info("WebSocket服务器初始化完成");
    }
    
    /**
     * 注入deviceService
     * 
     * @param deviceService 设备服务
     */
    public WebSocketServer(DeviceService deviceService, ObjectMapper objectMapper) {
        WebSocketServer.deviceService = deviceService;
        WebSocketServer.objectMapper = objectMapper;
    }
    
    /**
     * 连接建立
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("deviceId") String deviceId) {
        SESSIONS.put(deviceId, session);
        log.info("设备[{}]已连接", deviceId);
        
        // 更新设备在线状态
        if (deviceService != null) {
            deviceService.updateDeviceStatus(deviceId, 1);
        }
    }
    
    /**
     * 连接关闭
     */
    @OnClose
    public void onClose(@PathParam("deviceId") String deviceId) {
        SESSIONS.remove(deviceId);
        log.info("设备[{}]已断开连接", deviceId);
        
        // 更新设备离线状态
        if (deviceService != null) {
            deviceService.updateDeviceStatus(deviceId, 0);
        }
    }
    
    /**
     * 接收消息
     */
    @OnMessage
    public void onMessage(String message, @PathParam("deviceId") String deviceId) {
        log.info("收到设备[{}]的消息: {}", deviceId, message);
    }
    
    /**
     * 发生错误
     */
    @OnError
    public void onError(Throwable error, @PathParam("deviceId") String deviceId) {
        log.error("设备[{}]发生错误: {}", deviceId, error.getMessage());
    }
    
    /**
     * 向指定设备发送消息
     * 
     * @param deviceId 设备ID
     * @param message 消息
     */
    public void sendMessageToDevice(String deviceId, Object message) throws IOException {
        Session session = SESSIONS.get(deviceId);
        if (session != null && session.isOpen()) {
            String messageJson = objectMapper.writeValueAsString(message);
            session.getBasicRemote().sendText(messageJson);
            log.info("向设备[{}]发送消息: {}", deviceId, messageJson);
        } else {
            log.warn("设备[{}]不在线，无法发送消息", deviceId);
        }
    }
    
    /**
     * 向所有设备广播消息
     * 
     * @param message 消息
     */
    public void broadcastMessage(Object message) {
        String messageJson;
        try {
            messageJson = objectMapper.writeValueAsString(message);
        } catch (Exception e) {
            log.error("消息序列化失败: {}", e.getMessage());
            return;
        }
        
        SESSIONS.forEach((deviceId, session) -> {
            if (session.isOpen()) {
                try {
                    session.getBasicRemote().sendText(messageJson);
                } catch (IOException e) {
                    log.error("向设备[{}]发送消息失败: {}", deviceId, e.getMessage());
                }
            }
        });
        
        log.info("已向所有设备广播消息: {}", messageJson);
    }
    
    /**
     * 获取在线设备数量
     * 
     * @return 在线设备数量
     */
    public static int getOnlineCount() {
        return SESSIONS.size();
    }
} 